import pickle
from pathlib import Path
import numpy as np

from gist_extractor import image_GIST
from hog_extractor import image_HOG
from lbp_extractor import image_LBP


CUR_DIR = data_path = Path(__file__).resolve().parent


def scene15_gist_hog_lbp():
    """
    示例：提取Scene 15数据集的GIST、HOG、LBP，保存到文件
    """
    x = {
        'gist': list(),
        'hog': list(),
        'lbp': list()
    }
    y = []
    counter = 0
    for label_dir in CUR_DIR.joinpath('../dataset/15-Scene').iterdir():
        label = label_dir.name
        for img_path in label_dir.iterdir():
            counter += 1
            print('Loading image', counter, end='\r')
            gist = image_GIST(img_path)
            hog = image_HOG(img_path)
            lbp = image_LBP(img_path)
            x['gist'].append(gist)
            x['hog'].append(hog)
            x['lbp'].append(lbp)
            y.append(int(label))
    for i in x.keys():
        x[i] = np.array(x[i])
    y = np.array(y)

    # 保存文件
    print('\nDone. Save to local file.')
    pickle.dump([x, y], open(CUR_DIR / '../dataset/scene15-gist-hog-lbp.pkl', 'wb'))

    # 读取文件
    x, y = pickle.load(open(CUR_DIR / '../dataset/scene15-gist-hog-lbp.pkl', 'rb'))
    for i in x.keys():
        print('View', i, x[i].shape)
    print('label', y.shape)


if __name__ == '__main__':
    scene15_gist_hog_lbp()
